
module TableAccounting 
    (updateWeights,
     lookupWeight,
     emptyTable,
     permuteTry,
     permuteUnweighted,
     updateFunction,
     nprint,
     lookupM,
     RandTable)
where

{-
This module does the call accounting for the engine.  

Types:

ClauseID: This is an identifier for program clauses

BodyPredID: This is an identifier for body predicate clauses

GOAL is a special BPID

Stack: this looks like:

[(GOAL,CID1),
(PBID1,CID2),
(PBID2,CID3),
...

This is a list?

It is a stack of the decisions made thus far to get to the current program
state.  If the code is pure prolog, a selection rule could use this stack
to go directly to this point in the program.

  buildStack :: WAMData -> Stack
to interface between the WAM and the accounts

  callClause :: Stack -> Table -> [ClauseID] -> [ClauseID]
which given a current Stack and predicates potentially callable, gives an
optimized call order

something similar from randomize_neck

  updateTableForFailure :: Stack -> Table -> Double -> Table
given a stack and a table, calculates the new table based on the failure 
degree in the fourth argument

-}

import Prelude hiding (lookup)
import SParser
import RPTypes
import EngineTypes
import HistTree

import Data.Map (lookup,Map,empty,fromList,update,insert,showTree)
import Data.List hiding (lookup,insert)
import Maybe 
import Numeric
import System.Random

-- import Debug.Trace.Location
-- import Debug.Trace

type CallGraph = Map PredID [PredID]

emptyTable :: Int -> Double -> Double -> RandTable
emptyTable d l p = RT (HistTree.empty,d,l,p)

-- callGraph :: [Predicate] -> CallGraph
callGraph pl = fromList $ mergevals list 
    where list = foldl (++) [] [ pp $ clauselist pred | pred <- pl] 
          pp xs =  (foldl (++) [] [(pairs x) | x <- xs]) 
          pairs (CLA h b) = [(predID h,predID t) | t <- terms b]

-- XXX - no depth

lookupWeight :: Double -> ClauseIDRef -> [ClauseIDRef] -> RandTable -> RTEntry
lookupWeight avg nextcid history tab =
    let lookupval = foldl lookupk def [1 .. (length fullhist)]
	lookupk prev k = if isNothing w
			    then prev
			    else fromJust w
	    where w = HistTree.lookup klist (getRT tab)
		  klist = take k fullhist
	def = (RTE (fullhist,0,avg))
	fullhist = if depth == 0
		    then (nextcid:history)
		    else take depth (nextcid:history)
	depth = getDepth tab
    in scalePower lookupval (length fullhist) (getLambda tab)

scalePower :: RTEntry -> Int -> Double -> RTEntry
scalePower (RTE (list,c,weight)) newlen lambda = 
    (RTE (list,c,newweight))
    where newweight = weight * (lambda ** (fromIntegral ((length list) - newlen)))

lookupOneWeight :: Double -> ClauseIDRef -> [ClauseIDRef] -> RandTable -> RTEntry
lookupOneWeight avg nextcid history tab = 
    HistTree.findWithDefault (RTE (list,0,avg)) list table
    where list = if depth == 0
		    then (nextcid:history)
		    else take depth (nextcid:history)
	  prin = ("LOOKUPWEIGHT:  looking up: " ++ (show nextcid) ++ " and " ++ (show history) ++ 
                  "\n truncated to :" ++ (show list)++ 
		  " in: \n" ++ (HistTree.showTree nprint show 0 table) ++ "\n")
	  table = getRT tab
	  depth = getDepth tab


lookupM :: [ClauseIDRef] -> RandTable -> Maybe RTEntry
lookupM hist rt = HistTree.lookup hist tab
    where tab = getRT rt

-- switch - swithc first two elements

permuteTry:: [(WAMOpRaw, WAMOpRaw)] -> [WAMRef] -> RandTable -> [ClauseIDRef] -> [Double] -> Bool -> ([(WAMOpRaw, WAMOpRaw)],Maybe ClauseIDRef)
permuteTry instrs addrs tab history rands switch = 
    let cids = map getCIDAddrPair instrs
        defaultavg = 1
        lookupdefault entry = lookupWeight defaultavg (fst entry) history tab
        cidpairs = [ (cid,lookupdefault cid) | cid <- cids ]
        -- cidpairs = randompermute cidpairs0 (take (length cidpairs0) rands)
        -- this is a list of pairs [((CID,ADDR),(COUNT,WEIGHT))] which defaults to ((CID,ADDR),(0,1))
        nonzeros = filter (\(_,(RTE (_,x,_))) -> x /= 0) cidpairs
        avgnonzero = if (length nonzeros) >= 1
                        then (sum $ map (getWeight . snd) nonzeros) / (fromIntegral $ length nonzeros)
                        else 1.0
        startagain cid = lookupWeight avgnonzero (fst cid) [] tab  --- XXX this line is nonsense I think ...
        -- If we don't know the weight, assume it's the same as starting from scratch
	-- note ncidpairs are scaled by the permute scaler ...
        ncidpairs = [ (cid,(wrpScale scale) $ lookupWeight (getWeight $ startagain cid) (fst cid) history tab) | cid <- cids ]
        -- this is a list of pairs [((CID,ADDR),(COUNT,WEIGHT))] which defaults to ((CID,ADDR),(0,avgnonzero))
        -- where avgnonzero is the average of the non-zero elements
        wp0 = weightedRandomPermutation ncidpairs rands
	(wp,origelem) = switcher switch wp0 -- note origelem is a CR
        prin = ("***STARTS\nOriginal is: \n" ++ (printnice ncidpairs) ++ "\nUnmunged orig is: \n" ++ ( printnice cidpairs) ++ "\nRands are: " ++ (show rands) ++ "\nHistory is: " ++ (show history) ++ "\nWRP is : " ++ (show wp) ++ "\n***ENDS\n")
        table = getRT tab
	scale = getScale tab
    in  (zipWith3 reassignInstr (map (\(x,y) -> printCodeEntryInstAuxAlt x y) instrs) wp addrs,
	 maybe Nothing (Just . fst) origelem)

printnice l = foldl (\a b -> a ++ "\n" ++ b) "" (map pn1 l)

pn1((cid,_),RTE (list,count,weight)) = (show cid) ++ " : " ++ (showFFloat (Just 2) weight "") ++ " : " ++ (show count) ++ " --> " ++ (show (reverse list))

switcher False x = (x,Nothing)
switcher True (x:(y:r)) = ((y:(x:r)),Just x)
switcher True x = (x,Nothing)

wrpScale p (RTE (l,c,w)) = RTE (l,c,scaler w)
    where scaler x = x**p

getCIDAddrPair (x,y) = getCIDAddrPair1 $ printCodeEntryInstAuxAlt x y 

getCIDAddrPair1 (Try_there_alt ja _ _ cid)  = (cid,ja)
getCIDAddrPair1 (Retry_there_alt ja _ cid) = (cid,ja)
getCIDAddrPair1 (Trust_there_alt ja cid)    = (cid,ja)

reassignInstr (Try_there_alt _ size _ _) (cid,jumpaddr) next = wAMOpToWAMOpRawAlt (Try_there_alt jumpaddr size next cid)
reassignInstr (Retry_there_alt _ _ _)    (cid,jumpaddr) next = wAMOpToWAMOpRawAlt (Retry_there_alt jumpaddr next cid)   
reassignInstr (Trust_there_alt _ _)      (cid,jumpaddr) _    = wAMOpToWAMOpRawAlt (Randomize_trust_there_alt jumpaddr cid)   

-- Dodgy assign a random int to each op and sort technique
permuteUnweighted:: [(WAMOpRaw,WAMOpRaw)] -> [WAMRef] -> [Double] -> [(WAMOpRaw,WAMOpRaw)]
permuteUnweighted instrs addrs rands =
    let cids = map getCIDAddrPair instrs
        pairs = zip cids rands
        sorted = map fst (sortBy (\(_,x) (_,y) -> x `compare` y) pairs)
    in zipWith3 reassignInstr (map (\(x,y) -> printCodeEntryInstAuxAlt x y) instrs) sorted addrs


type CR = (ClauseIDRef,WAMRef)

-- Hill climbing version for debugging -- greatest first
wrpHillClimb :: [(CR,RTEntry)] -> [Double] -> [CR]
wrpHillClimb x _ = map (fst) (sortBy (\(_,RTE (_,_,x2)) (_,RTE (_,_,x1)) -> x1 `compare` x2 ) x)

weightedRandomPermutation ::[(CR,RTEntry)] -> [Double] -> [CR]
-- weightedRandomPermutation = wrp1
weightedRandomPermutation = wrpHillClimb


wrp1 :: [(CR,RTEntry)] -> [Double] -> [CR]
wrp1 [(cr,_)] _ = [cr]
wrp1 xs (r:rs) = 
    let sw = (sum $  (map weightFunction xs ))
        dart =  sw * r -- r should have range [0..1)
        (ele,newxs) = findWeight xs dart
    in ele:(wrp1 newxs rs)
    

weightFunction :: (CR,RTEntry) -> Double
-- weightFunction = snd . snd 
weightFunction (_,RTE (_,_,weight)) = weight

findWeight :: [(CR,RTEntry)] -> Double -> (CR,[(CR,RTEntry)]) 
findWeight [(acr,_)] _ = (acr,[])   
findWeight (cr@(acr,RTE (cl,cnt,ad)):rs) w 
    | (w - tw) < 0 = (acr,rs)
    | otherwise = (fst rest,(cr:(snd rest)))
    where rest = findWeight rs (w - ad)
          tw = weightFunction cr

{-
updateWeights :: RandTable -> Double -> Int -> [ClauseIDRef] -> Int -> RandTable
updateWeights = updateWeightAvg
-}

updateWeights :: RandTable -> RTEntry -> RandTable
updateWeights table rte =
    updateWeightChain table rte
	--in updateWeightInternal tab1 rte


updateWeightChain :: RandTable -> RTEntry -> RandTable
updateWeightChain tab (RTE (c,count,weight)) =
     foldl (\tl ins -> updateSingleWeight (updateFunction (updatescale ins)) tl (RTE (ins,count,weight))) tab 
            (map (\x -> drop x c) [0 .. ((length c)-1)])
     where rt = getRT tab
	   depth = getDepth tab
	   lambda = getLambda tab
	   updatescale it = (lambda**(fromIntegral ((length c) - (length it))))

updateSingleWeight :: (RTEntry -> RTEntry -> Maybe RTEntry) -> RandTable -> RTEntry  -> RandTable
updateSingleWeight uf tab rte@(RTE (c0,count,weight)) =
    if isJust $! (HistTree.lookup c rt)
       then RT ((HistTree.update $! (uf rte)) c rt,depth,lambda,scale)
       else RT (HistTree.insert c (fromJust $ uf (RTE (c,1,weight)) (RTE ([],0,0))) rt,depth,lambda,scale)
    where c = if depth /= 0
	         then take depth c0
                 else c0
	  rt = getRT tab
	  depth = getDepth tab
	  lambda = getLambda tab
	  scale = getScale tab
{-

updateWeightSucc :: RandTable -> [ClauseIDRef] -> [ClauseIDRef] -> Double -> Double  -> RandTable
updateWeightSucc tab oldh nh gamma val =
    let RTE (ohist,ocount,oval) = lookupWeight 1 (head oldh) (tail oldh) tab
	nval = gamma^^(length nh) * val
	rt = getRT tab
	depth = getDepth tab
	newweight = case val of 
                    v | (ohist == nh) && (val < oval) -> maxTree rt oldh -- val decreases
		      | nval > oval -> nval
		      | otherwise -> oval
	in RT(HistTree.update (\x -> Just x) oldh rt,depth)
-}

maxTree tree oldh = 
    HistTree.accumulateNodes ntree 0 maxer
    where maxer x cl1 (RTE (cl2,_,y)) = if cl1 == cl2
					   then max x y
					   else x
	  ntree = derefTree oldh tree

-- uf scale new old -> (maybe result)

updateFunction :: Double -> RTEntry -> RTEntry -> Maybe RTEntry
updateFunction = updateFunctionMax

updateFunctionDefault ::  Double -> RTEntry -> RTEntry -> Maybe RTEntry
updateFunctionDefault gamma (RTE(ins,count,weight)) _ = Just (RTE(ins,count,weight * gamma + weightAdd))

-- no +1 in next one ...
updateFunctionAvg :: Double -> RTEntry -> RTEntry -> Maybe RTEntry
updateFunctionAvg gamma (RTE(ins,ocount,weight0)) (RTE (cl,count,avg)) = 
    (Just $! (RTE (cl,
		  count + ocount,
		  ((avg * (fromIntegral count) + ((fromIntegral ocount) * weight * gamma))/(fromIntegral (count + ocount))))))
    where
    weight = weight0

updateFunctionMax :: Double -> RTEntry -> RTEntry -> Maybe RTEntry
updateFunctionMax gamma (RTE(ins,ocount,newweight)) (RTE (cl,count,oldweight)) = 
    (Just $! (RTE (cl,
                  count + ocount,
                  retweight)))
    where retweight = max oldweight (newweight*gamma + weightAdd)
	  prin = "Update; Oldweight is: " ++ (show oldweight) ++ " newweight is: " ++ (show newweight) ++ " gamma is: " ++ (show gamma)

mergevals [] = []
mergevals ((k,v):ks) = 
    let eqt (tk,_) = (k == tk)
        (goodvs,badvs) = Data.List.partition eqt ks
    in ((k,[snd v| v <- goodvs]):mergevals badvs)

terms :: CBody -> [Term]
terms CNIL = []
terms (CTERM t _) = [t]
terms (CCONJ c1 c2 _) = (terms c1) ++ (terms c2)
terms (CDISJ c1 c2 _) = (terms c1) ++ (terms c2)

nprint Nothing = (0,"")
nprint (Just (RTE (_,c,w))) =
    (w,(showFFloat (Just 2) w "") ++ " : " ++ (show c))
-- = if isNothing x then "" else (show $fromJust x) ++ "\n"

--- TESTING CODE
{-
permuteTester :: Int -> IO ()
permuteTester seed  =
    do putStrLn ("Rands: " ++ (show rands) ++ "\nInstrs: " ++ (show instrs) ++"\n")
       return ()
    where gen = mkStdGen seed
          (_,rands) = getR gen 2
          ii = [(Try_there 10 0 0 (CID "disj" 1 1)),(Trust_there 20 (CID "disj" 1 2))]
          addrs = [10,20]
          instrs = permuteTry ii addrs testrt [(CID "conj" 1 1)] rands

testrt =
    let e = emptyTable 2 0.8
        rt1 = updateWeights e rte1 
        rt2 = updateWeights rt1 rte2
    in rt2


rte1 = RTE ([(CID "disj" 1 1),(CID "conj" 1 1)],1,100)
rte2 = RTE ([(CID "disj" 1 2),(CID "conj" 1 1)],1,50)
-}

getR :: StdGen -> Int -> (StdGen,[Double])
getR rng k = (kth genaccum (rng,[]) k)

-- iterate kills the garbage collector
kth _ r 0 = r
kth f r n = kth f (f r) (n - 1)

genaccum :: (StdGen,[Double]) -> (StdGen,[Double])
genaccum (gen,ints) =
	let (i,g) = System.Random.random gen
	in (g,i:ints)

-- Printing good chains ...
{-
printChain :: RandTable -> [ClauseIDRef] -> [ClauseIDRef]
printChain rt start =
    let tab = 
-}